from jili.core.printlog import print
from jili.tool.convert import str2datetime
from gm import api
import datetime
ats2gm_mkcode={
    "SH":"SHSE",
    "SZ":"SZSE",
    # "BJ":"BJSE",
    # "HK":"HK",
    # "SHK":"SHK",
    # "ZHK":"ZHK",
    'CFFEX':"CFFEX",
    'CZCE':"CZCE",
    'DCE':"DCE",
    'GFEX':"GFEX",
    'INE':"INE",
    'SHFE' :"SHFE"
}
gm2ats_mkcode={v: k for k, v in ats2gm_mkcode.items()}
def obj2code(obj,mkcode=None):
    if not obj[0].isalpha():
        if mkcode is None:
            if "." in obj:
                obj,mkcode=obj.split(".")
                mkcode=ats2gm_mkcode[mkcode]
            else:
                if obj[0]>"3":
                    mkcode="SHSE"
                else:
                    mkcode="SZSE"
        else:
            if mkcode in ats2gm_mkcode.keys():
                mkcode=ats2gm_mkcode[mkcode]
    if mkcode is None:
        code= obj
    else:
        code=mkcode+"."+obj
    return code
def code2obj(code):
    mkcode,obj=code.split(".")
    if mkcode in gm2ats_mkcode.keys():
        mkcode=gm2ats_mkcode[mkcode]
    return obj,mkcode

periods={#60s, 300s, 900s, 1800s, 3600s
    "tick": "tick",
    "k1m":"60s",
    "k5m":"300s",
    "k15m":"900s",
    "k30m":"1800s",
    "k60m":"3600s",
    "k1d":"1d"
}
periods1={v: k for k, v in periods.items()}
ats2gm_fqtype={
    "qfq":1,
    "hfq":2,
    "bfq":0
}
class pricetool():
    def __init__(self,token):
        self.calls={}
        self.tool=api
        self.obj_infos={}
        self.tool.set_token(token)
    def run(self,strategy_id="",filename="",token="",serv_addr=""):
        api.run(strategy_id=strategy_id, filename=filename, mode=api.MODE_LIVE, token=token, serv_addr=serv_addr)
    def get_objinfo(self,obj):
        if obj not in self.obj_infos.keys():
            code=obj2code(obj)
            self.obj_infos[obj] = api.get_symbol_infos(sec_type1=1010, sec_type2=None, exchanges=None, symbols=code, df=False)
        return self.obj_infos[obj]
    def get_tradeday(self,count=1,end_date=None):
        if end_date is None:
            now = datetime.datetime.now()
            today = datetime.datetime(now.year, now.month, now.day)
        else:
            if isinstance(end_date,str):
                today=str2datetime(end_date)
            else:
                today=str2datetime(end_date)
        today=today+datetime.timedelta(days=1)
        tdays=api.get_previous_n_trading_dates(exchange="SHSE", date=today.strftime("%Y-%m-%d"), n=count)
        r=[]
        for i in tdays:
            r.append(str2datetime(i))
        return r
    def istradeday(self):
        now = datetime.datetime.now()
        today = datetime.datetime(now.year, now.month, now.day)
        r=self.get_tradeday(end_date=today)
        if today in r:
            return True
        else:
            return False
    def unsubscribe(self,obj, period='k30m'):
        code=obj2code(obj)
        period0=period[1:]
        if period[0]=="k":
            code0="bar_"
            code1=period0
        else:
            code0="tick_"
            code1=""
        key=code0+code+"_"+code1
        self.tool.unsubscribe(key)
    def subscribe(self,obj, peroid='k30m', callback=None):
        if isinstance(obj,str):
            code=[obj2code(obj)]
        else:
            code=[]
            for i in obj:
                code.append(obj2code(i))
        if peroid in periods.keys():
            period0=periods[peroid]
            print("subscribe",code,period0,obj,peroid)
            api.subscribe(symbols=code,frequency=period0,format="row")
            for i in code:
                if i not in self.calls.keys():
                    self.calls[i]={}
                self.calls[i][period0] = callback
        else:
            print("subscribe不支持的类型",peroid)
    def deal_tick(self,tick):
        code=tick["symbol"]
        obj,mkcode=code2obj(code)
        t={}
        t["dtype"] = "tick"
        t["obj"]=obj
        t["mkcode"]=mkcode
        t["timekey"] = tick["created_at"].replace(tzinfo=None)
        t["lastPrice"] = tick["price"]
        t["amount"] = tick["last_amount"]
        t["volume"]=tick["last_volume"]
        t["askPrice"]=[]
        t["askVol"]=[]
        t["bidPrice"] = []
        t["bidVol"] = []
        for i in tick["quotes"]:
            t["askPrice"].append(i["ask_p"])
            t["askVol"].append(i["ask_v"])
            t["bidPrice"].append(i["bid_p"])
            t["bidVol"].append(i["bid_v"])
        t["cumVolume"]=tick["cum_volume"]
        t["cumAmount"]=tick["cum_amount"]
        t["cumPosition"]=tick["cum_position"]
        t["code"]=code
        t["open"]=tick["open"]
        t["high"]=tick["high"]
        t["low"]=tick["low"]
        t["close"]=tick["price"]
        t["trade_type"]=tick["trade_type"]#1: '双开', 2: '双平', 3: '多开', 4: '空开', 5: '空平', 6: '多平', 7: '多换', 8: '空换'
        t["iopv"]=tick["iopv"]
        return t
    def deal_bar(self,bars):
        code = bars["symbol"]
        obj, mkcode = code2obj(code)
        bar = {}
        bar["obj"] = obj
        bar["mkcode"] = mkcode
        bar["timekey"] = bars["eob"].replace(tzinfo=None)
        frequency = bars["frequency"]
        bar["dtype"] = periods1[frequency]
        bar["open"] = bars["open"]
        bar["high"] = bars["high"]
        bar["low"] = bars["low"]
        bar["close"] = bars["close"]
        bar["volume"] = bars["volume"]
        bar["amount"] = bars["amount"]
        bar["position"] = bars["position"]
        bar["bob"] = bars["bob"]
        bar["eob"] = bars["eob"]
        return bar
    def on_tick(self,context, tick):
        # print('收到tick行情---', tick)
        t=self.deal_tick(tick)
        code=tick["symbol"]
        if code in self.calls.keys():
            if "tick" in self.calls[code].keys():
                self.calls[code]["tick"](t)
    def on_bar(self,context, barss):
        # print('收到bar行情---', barss)
        bars=barss[-1]
        code = bars["symbol"]
        bar=self.deal_bar(bars)
        frequency = bars["frequency"]
        if code in self.calls.keys():
            if frequency in self.calls[code].keys():
                self.calls[code][frequency](bar)
    def get_today_bar(self,obj,period="k1d",context=None):
        code=obj2code(obj)
        period=periods[period]
        bars=context.data(symbol=code, frequency=period, count=1)
        if bars:
            b=self.deal_bar(bars[-1])
            return b
        else:
            return {}
    def get_prebars(self,obj,count=1,period="k1d",fqtype="qfq",start_date='',end_date=''):
        code = obj2code(obj)
        period0 = periods[period]
        fqtype0=ats2gm_fqtype.get(fqtype,0)
        if end_date == '':
            end_date0 = self.get_tradeday(1)[0]
        else:
            end_date0 = end_date
        print("price_rq get_price", obj, start_date, end_date, end_date0, count, period, period0, fqtype, fqtype0)
        if count>0:
            bars=api.history_n(symbol=code, frequency=period0, count=count, end_time=end_date0, fields=None, skip_suspended=False,
          fill_missing=None, adjust=fqtype0, adjust_end_time='', df=False)

        else:
            bars=api.history(symbol=code, frequency=period0, start_time=start_date, end_time=end_date0, fields=None, skip_suspended=False,
          fill_missing=None, adjust=fqtype0, adjust_end_time='', df=False)
        bars0=[]
        if bars:
            for b in bars:
                b0=self.deal_bar(b)
                bars0.append(b0)
        return bars0

if __name__ == "__main__":
    import datetime
    pt=pricetool("ee2553d4a1e522d1cceda05efacd0c2ec0dcbf91")
    # print("get_objinfo",pt.get_objinfo("000001"))
    # print("get_objinfo",pt.get_objinfo("600000"))
    # print("get_objinfo",pt.get_objinfo("600000.SH"))
    # print("get_objinfo",pt.get_objinfo("SHSE.600000"))
    # print("istradeday",pt.istradeday())
    # print("get_tradeday", pt.get_tradeday(5,"2025-05-21"))
    # print("get_tradeday", pt.get_tradeday(5, "2025-05-18"))
    bars=pt.get_prebars("600000", 2, "k1d", "qfq", end_date= "2025-05-22")
    bars = pt.get_prebars("600000", 2, "k1d", "qfq", end_date="2025-05-22")
    bars = pt.get_prebars("600000", 2, "k1d", "qfq", end_date="2025-05-22")
    for i in bars:
        print("get_prebars", i["obj"],i["timekey"],i)